package br.unifor.g2cl.decorator;

import java.io.IOException;
import java.net.SocketAddress;
import java.util.HashMap;
import java.util.logging.Level;
import java.util.logging.Logger;

import net.sf.jgcs.Annotation;
import net.sf.jgcs.ClosedSessionException;
import net.sf.jgcs.ControlListener;
import net.sf.jgcs.ControlSession;
import net.sf.jgcs.GroupConfiguration;
import net.sf.jgcs.Message;
import net.sf.jgcs.Service;
import net.sf.jgcs.UnsupportedServiceException;
import net.sf.jgcs.utils.ControlSessionUtil;
import br.unifor.g2cl.DataSessionAdapter;
import br.unifor.g2cl.G2CLMessage;
import br.unifor.g2cl.IMarshalDataSession;
import br.unifor.g2cl.Util;

/**
 * The main goal of this service is to provide a way to fragment messages larger than a specified size before it is transmitted. 
 * 
 * On unreliable networks, retransmitting large messages may be very expensive, using this data session, only the lost fragment of the message need 
 * to be retransmitted.
 * <br/><br/>
 * This service is part of the Decorators group.
 * 
 * @author Henrique Teofilo (henriquetft@gmail.com)
 * @author Leandro Sales (leandro.shp@gmail.com)
 *
 */

/*
 *           1b        4b           4b            4b         ....
 * header: header message ID | message size | frag offset | payload
 */
public class FragDataSession extends DataSessionAdapter implements ControlListener{
	
	private final int EXPECTED_HEADER_LENGTH = 13;

	private static Logger logger = Logger.getLogger(FragDataSession.class
			.getName());

	private static final byte HEADER = 0xA;

	
	// XXX synchronize it?
	private HashMap<SocketAddress, HashMap<Integer, AsmMessage>> fragments;
	
	//private MembershipSession membershipSession;

	// XXX must be configurable
	private int fragSize = 3;

	// XXX need to be atomic int ???
	//private AtomicInteger myNextMessageId = new AtomicInteger(0);
	private int myNextMessageId;

	/**
	 * Size of the message fragments
	 * 
	 * @param fragSize
	 */
	public void setFragSize(String fragSize) {
		this.fragSize = Integer.parseInt(fragSize);
	}
	
	public int getNextMessageId() { 
		return ++myNextMessageId;
	}


	public FragDataSession(IMarshalDataSession data, ControlSession control) {		
		super(data);
		this.fragments = new HashMap<SocketAddress, HashMap<Integer, AsmMessage>>();
		ControlSessionUtil.redefineControlListener(control, this);
	}
	
	
	public int getFragSize() {
		return fragSize;
	}

	public void close() {
		data.close();
	}

	public Message createMessage() throws ClosedSessionException {
		return data.createMessage();
	}

	public GroupConfiguration getGroup() {
		return data.getGroup();
	}

	public void multicast(Message msg, Service service, Object cookie,
			Annotation... annotation) throws IOException,
			UnsupportedServiceException {
		FragmentGenerator fg = new FragmentGenerator((G2CLMessage) msg, fragSize);
		while (fg.hasNext()) {
			data.multicast(fg.next(), service, cookie, annotation);
		}
	}

	public void send(Message msg, Service service, Object cookie,
			SocketAddress destination, Annotation... annotation)
			throws IOException, UnsupportedServiceException {
		FragmentGenerator fg = new FragmentGenerator((G2CLMessage) msg, fragSize);
		while (fg.hasNext()) {
			data.send(fg.next(), service, cookie, destination, annotation); 
		}
	}

	public Object onMessage(Message msg) {
		G2CLMessage message = (G2CLMessage) msg;
		
		byte[] header = message.removeHeader();
		
		if (header.length == EXPECTED_HEADER_LENGTH && header[0] == HEADER) {
			byte[] payload = msg.getPayload();
			SocketAddress sender = msg.getSenderAddress();
			
			int msgId = Util.byteArrayToInt(header, 1);
			int msgSize = Util.byteArrayToInt(header, 5);
			int msgOffset = Util.byteArrayToInt(header, 9);
			
			if (logger.isLoggable(Level.FINEST)) {
				logger.finest("Fragmented message received. msgId: " + msgId +
						", msgSize: " + msgSize + ", msgOffset: " + msgOffset);
			}
			
			if (!fragments.containsKey(sender)) {
				fragments.put(sender, new HashMap<Integer, AsmMessage>());
			}
			HashMap<Integer, AsmMessage> peerFragTable = fragments.get(sender);
			
			if (!peerFragTable.containsKey(msgId)) {
				peerFragTable.put(msgId, new AsmMessage(msgSize));
			}
			
			AsmMessage asm = peerFragTable.get(msgId);
			asm.add(payload, msgOffset);
			
			if (asm.isComplete()) {
				try {
					if (logger.isLoggable(Level.FINER)) {
						logger.finer("Message id " + msgId + " re-assembled. "+
								"Delivering it...");
					}
					G2CLMessage bigMessage = (G2CLMessage) data.createMessage();
					bigMessage.setPayload(asm.getMessage());
					peerFragTable.remove(msgId);
					bigMessage.setSenderAddress(sender);
					return realMessageListener.onMessage(bigMessage);
				} catch (ClosedSessionException e) {
					logger.log(Level.WARNING, "received message and then " +
							"disconnected? (the last message was not delivered)", e);
				}
			}
			
			
		} else {
			if (logger.isLoggable(Level.SEVERE)) {
			logger.severe("ERROR! Discarding message. This is not a " +
					"fragmented message. Make sure that the stack of " +
					"datasessions are configured same way at the two " +
					"endpoints.");
			}
		}
		
		return null;
	}

	/** Iterator-like class to generate fragments of a message */
	 private class FragmentGenerator {

		private int fragSize;
		private byte[] payload;
		private int total;
		private int p;
		private G2CLMessage message;
		private int messageId;

		public FragmentGenerator(G2CLMessage message, int fragSize) {
			this.message = message;
			this.fragSize = fragSize;
			this.payload = message.getPayload();
			this.total =  payload.length;
			this.p = 0;
			this.messageId = getNextMessageId();
			
			if (total > fragSize && logger.isLoggable(Level.FINER)) {
				logger.finer("message id " + messageId + " will be fragmented");
			}
			
		}

		public boolean hasNext() {
			return p < total;
		}

		public G2CLMessage next() throws ClosedSessionException {
			// size of the next fragment
			int size = Math.min(fragSize, total - p);

			byte[] frag = new byte[size];
			System.arraycopy(payload, p, frag, 0, frag.length);

			// filling header
			byte[] header = new byte[13];
			header[0] = HEADER;
			Util.intToByteArray(header, 1, messageId);
			Util.intToByteArray(header, 5, total);
			Util.intToByteArray(header, 9, p);
			
			if (logger.isLoggable(Level.FINEST)) {
				String s = "generating frag: id({0}) total({1}) offset({2})";
				logger.log(Level.FINEST, s, new Object[] { messageId,
						total, p });
			}

			G2CLMessage fragMsg = (G2CLMessage) data.createMessage();
			fragMsg.setSenderAddress(message.getSenderAddress());
			fragMsg.setPayload(frag);
			fragMsg.addHeader(header);
			p += frag.length;
			return fragMsg;
		}
	}


	/* ***************************** */
	/*        Control Listener       */ 
	 /* ***************************** */ 
	public void onFailed(SocketAddress failedMember) {
		fragments.remove(failedMember);
	}

	public void onJoin(SocketAddress joinedMember) {
		//Nothing to do
	}

	public void onLeave(SocketAddress leavedMember) {
		fragments.remove(leavedMember);
	}
}

/** Class to help on re-assembling of messages */
class AsmMessage {
	private static Logger logger = Logger.getLogger(AsmMessage.class
			.getName());
	private byte[] message;
	
	private int countBytes;
	
	public AsmMessage(int size) {
		message = new byte[size];
		countBytes = 0;
	}
	
	public void add(byte[] fragment, int offset) {
		if (logger.isLoggable(Level.FINEST)) {
			logger.finest("==============\n" +
			"original message size: " + message.length + "\n" +
			"bytes received: " + countBytes + "\n" +
			"frag added size: " + fragment.length + "\n" +
			"frag offset: " + offset + "\n" +
			"================");
		}
		countBytes += fragment.length;
		System.arraycopy(fragment, 0, message, offset, fragment.length);
	}
	
	public boolean isComplete() {
		return countBytes == message.length;
	}
	
	public byte[] getMessage() {
		return message;
	}
}
